from const import *
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import expit

class NewTon3:
    def __init__(self, A, b):
        self.A = A
        self.n = A.shape[1]
        self.p = A.shape[0]
        self.v0 = np.random.random(self.p)
        self.b = b

    def f(self, x):
        return x @ np.log(x)

    # max -b^Tv-sum(-exp^(A^Tv[i]-1))
    def gdual(self, v):
        return -self.gconvex(v)

    # min b^Tv+sum(-exp^(A^Tv[i]-1))
    def gconvex(self, v):
        sum = np.dot(self.b,v)
        fstar = np.exp(- (np.transpose(self.A) @ v) - 1)
        sum += np.sum(fstar)
        return sum
    
    def gradient(self, v):
        fstar = np.exp(- (np.transpose(self.A) @ v) - 1)
        return self.b - self.A @ fstar

    def hessian(self, v):
        res = np.zeros((len(v), len(v)), dtype=np.float64)
        fstar = np.exp(- (np.transpose(self.A) @ v) - 1)
        for i in range(len(v)):
            for j in range(len(v)):
                res[i][j] = (self.A[i] * self.A[j]) @ fstar
        return res
    
    def dv(self, v):
        return - np.linalg.inv(self.hessian(v)) @ self.gradient(v)
    
    def declamda(self, v, dv):
        return dv @ self.hessian(v) @ dv
    
    def search(self):
        v0 = self.v0
        ls_f = [self.gdual(v0)]
        while True:
            t = 1
            dv = self.dv(v0)
            dec = 0.5 * self.declamda(v0, dv)
            if dec < epsilon:
                break
            v1 = v0 + t * dv
            while self.gconvex(v1) > self.gconvex(v0) - alpha * t * dec:
                t = beta * t
                v1 = v0 + t * dv
            v0 = v1
            ls_f.append(self.gdual(v0))
            print(f"epoach :{len(ls_f) - 1}  λ(x)^2: {dec}  g(v): {ls_f[-1]}")
        return v0, ls_f
    
    def calx(self, v):
        return np.exp(- np.transpose(self.A) @ v - 1)
    
    def search_plot(self):
        v, ls = self.search()
        print(f"x={self.calx(v)}")
        print(f"v={v}")
        print(f"g(v) = {ls[-1]}  f(x) = {self.f(self.calx(v))}")
        plt.plot(range(len(ls)),ls)
        plt.xlabel("epoach")
        plt.ylabel("g(v)")
        plt.show()

if __name__ == "__main__":
    A = np.loadtxt("A.txt")
    b = np.loadtxt("b.txt")
    new = NewTon3(A, b)
    new.search_plot()
